#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
FLUX训练数据准备脚本
用于组织和验证训练数据集
"""

import os
import shutil
import argparse
from pathlib import Path
import json

def create_directory_structure(base_dir):
    """创建训练数据目录结构"""
    directories = [
        "training_data_all/images",
        "training_data_all/output",
        "training_data_all/reg_images",  # 可选的正则化图像
    ]
    
    for dir_path in directories:
        full_path = Path(base_dir) / dir_path
        full_path.mkdir(parents=True, exist_ok=True)
        print(f"✓ 创建目录: {full_path}")

def validate_image_files(image_dir):
    """验证图像文件"""
    image_extensions = {'.jpg', '.jpeg', '.png', '.webp', '.bmp'}
    image_files = []
    
    if not os.path.exists(image_dir):
        print(f"❌ 图像目录不存在: {image_dir}")
        return image_files
    
    for file_path in Path(image_dir).rglob('*'):
        if file_path.suffix.lower() in image_extensions:
            image_files.append(file_path)
    
    print(f"✓ 找到 {len(image_files)} 个图像文件")
    return image_files

def validate_caption_files(image_dir, caption_ext='.txt'):
    """验证标题文件"""
    caption_files = []
    missing_captions = []
    
    image_extensions = {'.jpg', '.jpeg', '.png', '.webp', '.bmp'}
    
    for file_path in Path(image_dir).rglob('*'):
        if file_path.suffix.lower() in image_extensions:
            caption_path = file_path.with_suffix(caption_ext)
            if caption_path.exists():
                caption_files.append(caption_path)
            else:
                missing_captions.append(file_path)
    
    print(f"✓ 找到 {len(caption_files)} 个标题文件")
    if missing_captions:
        print(f"⚠️  缺少 {len(missing_captions)} 个标题文件")
        for img_path in missing_captions[:5]:  # 只显示前5个
            print(f"   缺少: {img_path.name}")
        if len(missing_captions) > 5:
            print(f"   ... 还有 {len(missing_captions) - 5} 个文件")
    
    return caption_files, missing_captions

def create_sample_captions(image_dir, caption_ext='.txt'):
    """为缺少标题的图像创建示例标题"""
    image_extensions = {'.jpg', '.jpeg', '.png', '.webp', '.bmp'}
    created_count = 0
    
    for file_path in Path(image_dir).rglob('*'):
        if file_path.suffix.lower() in image_extensions:
            caption_path = file_path.with_suffix(caption_ext)
            if not caption_path.exists():
                # 创建示例标题
                sample_caption = f"masterpiece, best quality, {file_path.stem}"
                with open(caption_path, 'w', encoding='utf-8') as f:
                    f.write(sample_caption)
                created_count += 1
    
    if created_count > 0:
        print(f"✓ 创建了 {created_count} 个示例标题文件")
        print("⚠️  请手动编辑这些标题文件以提供准确的描述")

def copy_files_to_training_dir(source_dir, target_dir, file_patterns=None):
    """复制文件到训练目录"""
    if file_patterns is None:
        file_patterns = ['*.jpg', '*.jpeg', '*.png', '*.webp', '*.bmp']
    
    target_path = Path(target_dir)
    target_path.mkdir(parents=True, exist_ok=True)
    
    copied_count = 0
    for pattern in file_patterns:
        for file_path in Path(source_dir).glob(pattern):
            target_file = target_path / file_path.name
            if not target_file.exists():
                shutil.copy2(file_path, target_file)
                copied_count += 1
    
    print(f"✓ 复制了 {copied_count} 个文件到 {target_dir}")

def create_dataset_info(base_dir):
    """创建数据集信息文件"""
    info = {
        "dataset_name": "FLUX训练数据集",
        "created_date": str(Path().cwd()),
        "directories": {
            "images": "training_data_all/images",
            "output": "training_data_all/output",
            "reg_images": "training_data_all/reg_images"
        },
        "config_files": {
            "lora_training": "dataset_1024_bs2.toml",
            "finetune_training": "dataset_1024_bs1.toml"
        },
        "notes": [
            "图像文件支持: jpg, jpeg, png, webp, bmp",
            "标题文件扩展名: .txt",
            "建议图像分辨率: 512x512 到 2048x2048",
            "标题文件应与图像文件同名（除了扩展名）"
        ]
    }
    
    info_file = Path(base_dir) / "dataset_info.json"
    with open(info_file, 'w', encoding='utf-8') as f:
        json.dump(info, f, indent=2, ensure_ascii=False)
    
    print(f"✓ 创建数据集信息文件: {info_file}")

def main():
    parser = argparse.ArgumentParser(description="FLUX训练数据准备工具")
    parser.add_argument("--base-dir", default=".", help="基础目录路径")
    parser.add_argument("--source-dir", help="源图像目录路径")
    parser.add_argument("--create-captions", action="store_true", help="为缺少标题的图像创建示例标题")
    parser.add_argument("--copy-images", action="store_true", help="复制图像到训练目录")
    
    args = parser.parse_args()
    
    print("========================================")
    print("FLUX训练数据准备工具")
    print("========================================")
    
    base_dir = Path(args.base_dir)
    
    # 1. 创建目录结构
    print("\n1. 创建目录结构...")
    create_directory_structure(base_dir)
    
    # 2. 验证现有数据
    images_dir = base_dir / "training_data_all" / "images"
    print(f"\n2. 验证训练数据...")
    print(f"图像目录: {images_dir}")
    
    if images_dir.exists():
        image_files = validate_image_files(images_dir)
        caption_files, missing_captions = validate_caption_files(images_dir)
        
        if missing_captions and args.create_captions:
            print(f"\n3. 创建示例标题...")
            create_sample_captions(images_dir)
    else:
        print(f"⚠️  图像目录不存在: {images_dir}")
    
    # 3. 复制源图像（如果指定）
    if args.source_dir and args.copy_images:
        print(f"\n3. 复制图像文件...")
        source_dir = Path(args.source_dir)
        if source_dir.exists():
            copy_files_to_training_dir(source_dir, images_dir)
        else:
            print(f"❌ 源目录不存在: {source_dir}")
    
    # 4. 创建数据集信息
    print(f"\n4. 创建数据集信息...")
    create_dataset_info(base_dir)
    
    # 5. 显示使用说明
    print(f"\n========================================")
    print("数据集准备完成！")
    print("========================================")
    print("下一步操作:")
    print("1. 将训练图像放入: training_data_all/images/")
    print("2. 为每个图像创建对应的.txt标题文件")
    print("3. 编辑标题文件，提供准确的图像描述")
    print("4. 运行FLUX训练脚本:")
    print("   - LoRA训练: ./flux_24g_training.sh")
    print("   - 高级LoRA: ./flux_24g_optimized.sh")
    print("   - 微调训练: ./flux_24g_finetune.sh")
    print("\n标题文件示例:")
    print("masterpiece, best quality, (1girl), in white shirts, upper body, looking at viewer, simple background --n low quality, worst quality, bad anatomy,bad composition, poor, low effort")

if __name__ == "__main__":
    main()
